"""
author: liuxu
date: 2025/7/7
description: 获取硬件环境信息（包括dcu，cpu，内存，硬盘等信息）
"""
import psutil
import platform
import threading
import traceback
import re
import os
import sys

from common.ssh_command import exec_cmd
from common.logger_ctrl import mylogger


class SystemHwInfo:
    """硬件信息单例类，全局只初始化一次"""
    _instance = None
    _lock = threading.Lock()

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self):
        self._sys_hw_info = None

    @property
    def sys_hw_info(self):
        if self._sys_hw_info is None:
            self._sys_hw_info = self.get_dcu_info()

        return self._sys_hw_info

    @staticmethod
    def get_cpu_info():
        cpu_info = {
            "cores": psutil.cpu_count(logical=False),
            "logical_cores": psutil.cpu_count(logical=True),
            "usage": psutil.cpu_percent(interval=1)
        }
        return cpu_info

    @staticmethod
    def get_memory_info():
        mem = psutil.virtual_memory()
        memory_info = {
            "total": mem.total,
            "available": mem.available,
            "used": mem.used,
            "percent": mem.percent
        }
        return memory_info

    @staticmethod
    def get_sys_info():
        os_info = {
            "system": platform.system(),
            "release": platform.release(),
            "version": platform.version()
        }
        return os_info

    def get_dcu_info(self):
        """
        获取实际可用dcu卡信息
        :return:
        """
        all_dcu_detail_l = self._get_dcu_info_list_from_rocminfo()

        # 获取并解析HIP_VISIBLE_DEVICES环境变量
        hip_visible = os.getenv('HIP_VISIBLE_DEVICES', '')
        hip_visible_valid_indices = [
            int(idx) for idx in hip_visible.split(',')
            if idx.strip().isdigit() and 0 <= int(idx) < len(all_dcu_detail_l)
        ]
        valid_device_id_l = hip_visible_valid_indices or list(range(len(all_dcu_detail_l)))  # 无有效环境变量则使用rocminfo获取的设备列表
        # 根据是否有有效索引决定使用哪些设备
        available_dcu_detail_l = [all_dcu_detail_l[idx] for idx in valid_device_id_l]

        # 构建返回数据结构
        dcu_info_d = {
            "count": len(available_dcu_detail_l),  # 卡个数
            "devices_id": valid_device_id_l,
            "gfx_arch": available_dcu_detail_l[0].get("name") if available_dcu_detail_l else "N/A",
            "avail_dcu_detail_l": available_dcu_detail_l,
            "dcu_type": self._get_dcu_type_from_hysmi()
        }
        mylogger.info("Available dcu detail info: {}".format(dcu_info_d))
        return dcu_info_d

    @staticmethod
    def _get_dcu_type_from_hysmi():
        """获取dcu卡类型"""
        cmd = "hy-smi --showproductname"
        try:
            res, output = exec_cmd(cmd)
            if not res:
                raise RuntimeError(f"{cmd} command failed")
            return re.search(r"Card Series:\s+(\w+\s?\d+)", output).group(1)
        except Exception as e:
            mylogger.error(f"Failed to run {cmd} command:\n {e}\n{traceback.format_exc()}")
            sys.exit(-1)

    @staticmethod
    def _get_dcu_info_list_from_rocminfo():
            try:
                res, output = exec_cmd("rocminfo")
                if not res:
                    raise RuntimeError("rocminfo command failed")
            except Exception as e:
                mylogger.error(f"Failed to run rocminfo command:\n {e}\n{traceback.format_exc()}")
                sys.exit(-1)

            # """解析rocminfo输出文本"""
            gpu_info_l = []

            # 分割为每个Agent的信息块
            agent_block_result = re.split(r'\n\*{7,}\s*\nAgent\s+\d+\s*\n\*{7,}\s*\n', output)[1:]
            for block_info in agent_block_result:
                def get_info(pattern, text=block_info, flags=0):
                    """辅助函数：提取匹配内容或返回N/A"""
                    match = re.search(pattern, text, flags)
                    return match.group(1).strip() if match else 'N/A'

                # 提取gpu信息，忽略cpu信息收集
                if get_info(r'Device Type:\s+(\w+)') == "CPU":
                    continue

                # 提取GPU信息
                info = {
                    'name': get_info(r'Name:\s+(.*?)(?:\n|$)'),
                    'uuid': get_info(r'Uuid:\s+(.*?)(?:\n|$)'),
                    'marketing_name': get_info(r'Marketing Name:\s+(.*?)(?:\n|$)'),
                    'vendor_name': get_info(r'Vendor Name:\s+(.*?)(?:\n|$)'),
                }

                gpu_info_l.append(info)
            return gpu_info_l



if __name__ == '__main__':
    res=SystemHwInfo().get_dcu_info()
    mylogger.debug(res)




# # 在conftest.py中初始化
# @pytest.fixture(scope="session", autouse=True)
# def init_system_info():
#     """会话级fixture，自动初始化系统信息单例"""
#     sys_info = SystemHwInfo()
#     sys_info.print_summary()
#     return sys_info
#
#
# # 测试基类
# class BaseTest:
#     """测试基类，提供访问系统信息的方法"""
#
#     @pytest.fixture(autouse=True)
#     def inject_system_info(self, init_system_info):
#         """注入系统信息到测试实例"""
#         self.sys_info = init_system_info
#
#     def get_available_gpus(self):
#         """获取可用GPU数量（公共方法示例）"""
#         return self.sys_info.gpu_info["gpu_count"]
#
#     def get_cpu_cores(self):
#         """获取CPU核心数"""
#         return self.sys_info.cpu_info["cores"]
#
#     def get_total_memory(self):
#         """获取总内存（GB）"""
#         return self.sys_info.memory_info["total"] / 1024 ** 3